public class AVLTree {
    static class AVLNode{
        int key;
        Object value;
        AVLNode left;
        AVLNode right;
        int height = 1; //高度

        public AVLNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(int key) {
            this.key = key;
        }

        public AVLNode(int key, Object value, AVLNode left, AVLNode right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }
    }
    //求节点高度
    private int height(AVLNode node){
        return node == null ? 0 : node.height;
    }
    //更新节点高度
    private void updateHeight(AVLNode node){
        node.height = Integer.max(height(node.left),height(node.right));
    }
    //平衡因子 = 左子树高度 - 右子树高度
    private int bf(AVLNode node){
        return height(node.left) - height(node.right);
    }
    //旋转
    //右旋树
    private AVLNode rightRotate(AVLNode red){
        AVLNode yellow = red.left;
        AVLNode green = yellow.right;
        yellow.right = red;
        red.left = green;
        updateHeight(red);
        updateHeight(yellow);
        return yellow;
    }
    //左旋树
    private AVLNode leftRotate(AVLNode red){
        AVLNode yellow = red.right;
        AVLNode green = yellow.left;
        yellow.left = red;
        red.right = green;
        updateHeight(red);
        updateHeight(yellow);
        return yellow;
    }
    //左右旋树
    //先旋左子树，再旋根节点
    private AVLNode leftRightRotate(AVLNode node){
        node.left = leftRotate(node.left);
        return rightRotate(node);
    }
    //右左旋树
    //先旋右子树，
    private AVLNode rightLeftRotate(AVLNode node){
        node.right = rightRotate(node.right);
        return leftRightRotate(node);
    }
    //检查节点是否失衡，重新平衡代码
    private AVLNode balance(AVLNode node){
        if(node == null){
            return null;
        }
        int bf = bf(node);
        if(bf > 1 && bf(node.left) >= 0){
            return leftRotate(node);
        }else if(bf > 1 && bf(node.left) < 0){
            return leftRightRotate(node);
        }else if(bf < 1 && bf(node.right) >= 0){
            return rightLeftRotate(node);
        }else if(bf < 1 && bf(node.right) < 0){
            return rightRotate(node);
        }
        return node;
    }
    AVLNode root;
    private void put(int key,Object value){
        root = doPut(root,key,value);
    }
    private AVLNode doPut(AVLNode node,int key,Object value){
        //找到空位，创建新节点
        if(node == null){
            return new AVLNode(key,value);
        }
        //已存在，更新
        if(key == node.key){
            node.value = value;
            return node;
        }
        //继续查找
        if(key < node.key){
            node.left = doPut(node.left,key,value);
        }else{
            node.right = doPut(node.right,key,value);
        }
        updateHeight(node);
        return balance(node);
    }
    public void remove(int key){
        doRemove(root,key);
    }
    private AVLNode doRemove(AVLNode node,int key){
        //1.node == null
        if(node == null){
            return null;
        }
         //2.没找到key
        if(key < node.key){
            node.left = doRemove(node.left,key);
        }else if(node.key < key){
            node.right = doRemove(node.right,key);
        }else{
            //3.找到key 1）没有孩子 2）只有一个孩子 3）有两个孩子
            if(node.left == null && node.right == null){
                return null;
            }else if(node.left == null){
                return node.right;
            }else if(node.right == null){
                return node.left;
            }else{
                AVLNode s = node.right;
                while(s.left != null){
                    s = s.left;
                }
                // s 后继结点
                s.right = doRemove(node.right,s.key);
                s.left = node.left;
                node = s;
            }
        }
        //4.更新高度
         updateHeight(node);
        //5. balance
        return balance(node);
    }
}
